import os
import sys
from jsonpath import jsonpath
from locust import HttpUser, task, LoadTestShape, constant

curPath = os.path.abspath(os.path.dirname(__file__))
rootPath = os.path.split(curPath)[0]
PathProject = os.path.split(rootPath)[0]
sys.path.append(rootPath)
sys.path.append(PathProject)


class QuickStartUser(HttpUser):
    host = "http://127.0.0.1:8000"
    wait_time = constant(1)

    @task
    def api(self):
        with self.client.get(url="/locust_test", catch_response=True, name='测试') as response:
            print(response.json())
            if response.status_code == 200:
                try:
                    assert jsonpath(response.json(), '$..code')[0] == 1
                except:
                    response.failure("Timeout")
            else:
                response.failure(response.status_code)


class CustomShape(LoadTestShape):
    # 设置时限
    time_limit = 600
    # 设置产生率
    spawn_rate = 20

    def tick(self):
        """
        设置 tick()函数
        并在tick()里面调用 get_run_time()方法
        """

        # 调用get_run_time()方法
        run_time = self.get_run_time()
        print(f'当前运行时间：{run_time}')
        # 运行时间在 10分钟之内，则继续执行
        if run_time < self.time_limit:
            user_count = round(run_time, -2)
            print(f'用户数：{user_count}')
            # 返回user_count,spawn_rate这两个参数
            tick_data = (user_count, self.spawn_rate)
            return tick_data
        return None


if __name__ == '__main__':
    os.system("locust -f 负载图_时间峰值策略.py")
